feat(refit): add ModelExpress Megatron-to-vLLM reshard path - #3632

What this PR does

This PR adds the complete NeMo-RL side of the ModelExpress mx_reshard refit path:

Each part has a narrow job in making trainer-to-generator refits usable from NeMo-RL:

• Describe what every Megatron rank owns. Megatron stores model weights already split across TP, EP, and ETP ranks. The publisher records each rank's real piece and its position in the full tensor, so ModelExpress can read it directly instead of first gathering a complete model on every trainer.
• Translate Megatron names into the names vLLM expects. The trainer and generator use different names for the same weights, and fused tensors may map to several HF tensors. Megatron-Bridge supplies that relationship; this adapter preserves the required Q/K/V and gate/up order and fails when a name is ambiguous.
• Connect ModelExpress to NeMo-RL's normal refit lifecycle. The WeightSynchronizer now starts publishers and receivers, waits for all required ranks, runs the transfer, commits the new version, and shuts both sides down in a safe order. Callers do not need a separate ModelExpress-only control path.
• Prevent partial models from becoming active. A refit proceeds only after every expected publisher is ready, and one receiver failure prevents the version from being committed. This avoids serving a model assembled from a mixture of old and new weights.
• Make correctness and delays visible. Parameter verification catches transfers that report success but install unchanged or wrong bytes. Separate stage records show whether time was spent waiting for publishers, moving bytes, or installing tensors, which keeps performance investigations from guessing.
• Provide a functional GRPO gate that cannot pass by accident. The test checks that the recipe launches, the model actually changes, and trainer/generator probabilities remain aligned. Portable recipes and analysis tools keep the same checks reproducible without upstreaming private cluster names or images.
• Publish global QKV geometry from each live layer. Nemotron-style layers can have fewer KV heads than TP ranks, and different layers may use different attention shapes. Reading each owning linear_qkv module avoids 2 // 8 = 0 and prevents one root-model setting from being incorrectly reused for every layer.

ModelExpress dependency

This path needs ModelExpress main at or after f4660d2. ModelExpress #635 carried the transport work and the global QKV interval alias builder this PR's QKV contract requires, and it merged on 2026-08-19, so no ModelExpress branch is needed any more. It was squash-merged, so pin by commit rather than by ancestry.

Neither half is useful alone: NeMo-RL publishes the geometry, ModelExpress maps the intervals. ModelExpress #482 does not work — it lacks VllmReshardReceiver (modelexpress.engines.vllm.refit.receiver), which is what this PR imports.

Enable the path with policy.generation.refit_transport: mx_reshard. This is not the legacy cluster.weight_sync.method: mx path; the two are unrelated.

Architecture

flowchart LR
    M[Live Megatron layer] --> G[Per-layer Q/KV geometry resolver]
    G --> C[Parameter classifier]
    C --> B[Megatron-Bridge name map]
    B --> P[ModelExpress publisher]
    P --> Q[Publisher quorum]
    Q --> R[vLLM reshard receiver]
    R --> I[Install and verify]
    I --> V[Commit one version]

    W[WeightSynchronizer] --> P
    W --> Q
    W --> R
    W --> V
    W --> S[Ordered receiver/publisher shutdown]

For Q=64, KV=2 and trainer TP8, NeMo-RL now publishes global num_heads, num_kv_heads, head_dim, and qkv_interleave from each live linear_qkv.config. It does not publish num_kv_heads_local=0, and it does not pretend every rank owns a replicated KV head. ModelExpress receives each rank's real fused-row range and performs the interval mapping.

Why per-layer geometry matters

Nemotron-family models may use different attention geometry in different layers. Reading one root model config would stamp the same Q/KV shape onto every QKV tensor. The resolver therefore walks to the owning live linear_qkv module first and uses the root values only as a compatibility fallback. The chosen geometry must match the tensor's global fused-row count or publication fails closed.

Performance measured on real model tensors

These are three different experiments, not interchangeable rows from one run:

• Dense performance A/B — Qwen3-4B-Thinking-2507. The Megatron trainer uses 16 GPUs as TP2 x PP1 x DP8. The generator uses the other 16 GPUs as four independent vLLM replicas, each TP4. Every model weight therefore exists on eight trainer DP replicas, while each generation replica needs one TP4-sharded copy. This is the setup that measures replica deduplication: the baseline reads all eight byte-identical trainer replicas; the new path identifies them by geometry and digest and reads one.
• MoE coverage and timing — Qwen3-30B-A3B-Instruct-2507. This is a 48-layer, 128-expert MoE model. The trainer uses 16 GPUs as TP2 x PP1 x EP4 x ETP1 x DP2; generation again uses 16 GPUs as four vLLM TP4 replicas. This setup exercises grouped-expert naming and fused expert capture, then reshards TP2/EP4 trainer ownership into TP4 generator ownership. The model has Q=32, KV=4, head_dim=128, so it exercises the global QKV path but not KV heads below TP.
• KV-heads-below-TP correctness — reduced Qwen3, Q=64/KV=2/head_dim=128. One Megatron TP8 model refits one vLLM TP8 model, 8 trainer + 8 generator GPUs. Because each node has four GPUs, both TP8 groups span two nodes. This arm asks whether sparse Q/K/V row ownership is mapped bit-for-bit; it is deliberately not a performance benchmark.

The dense and MoE runs use 8 nodes x 4 NVIDIA GB200 = 32 GPUs, split into 16 non-colocated trainer GPUs and 16 generation GPUs. Nodes are arm64; each has four 189 GB GPUs and four RDMA rails. Weight refit uses NIXL/UCX RDMA. NCCL is used only for Megatron-internal collectives, not as the refit transport. All three arms use BF16.

For performance rows, the methodology is three independent cold starts per arm, then one warm-up plus ten measured warm refits. "Fleet-critical receiver latency" means: for each refit, take the slowest receiver rank because the fleet cannot continue until that rank finishes; then report the distribution of those per-refit maxima. It is not an average-rank latency.

An attribution note, so these numbers are read correctly: the byte and latency reductions below come from the ModelExpress transport changes (replica deduplication, descriptor caching, batched install, fused expert capture). This PR is the NeMo-RL half that makes that transport reachable from GRPO, enforces the version/quorum ordering it needs, and supplies the telemetry that produced every figure here. The rows are what a NeMo-RL user actually gets end to end, not a claim that this diff alone caused them.

Dense Qwen3-4B: replica-deduplication A/B

This table belongs only to the dense setup above: Qwen3-4B, Megatron TP2/DP8 on 16 GPUs -> four vLLM TP4 replicas on 16 GPUs. Same model checkpoint, prompts, NeMo-RL tree, nodes and placement in both arms; only ModelExpress differs (9abaf2dd, the transport's parent, versus 3ddb9b16, the six-change candidate).

Comparing median-of-three cold starts against median-of-three, the receiver speedup is 7.57x (4.526 s -> 0.598 s). Byte counts were bit-identical across all three cold starts in each arm, with fallback and full_pull_sources at 0 throughout.

Please read the throughput row carefully and do not quote it as a regression. Per-rank throughput goes down while wall clock improves 7.6x. The baseline sustains a higher rate precisely because it moves 8x more data with eight concurrent sources per shard, which amortises per-descriptor overhead well; the new path finishes so much sooner that fixed costs occupy a larger share of a much shorter transfer. The win is not moving redundant bytes at all — the trainer runs DP8, so every weight exists on eight data-parallel replicas and the receiver was re-reading the same bytes — rather than driving the fabric harder.

The framework-visible 5.84x is the honest user-facing number, and it is smaller than the receiver-side ratio because it also carries publisher-side preparation and Ray call overhead.

MoE Qwen3-30B-A3B: expert coverage

This table belongs only to the MoE setup above: Qwen3-30B-A3B, Megatron TP2/EP4/DP2 on 16 GPUs -> four vLLM TP4 replicas on 16 GPUs. "Before" is the same model and topology without fused grouped-expert capture; "after" is the current reshard path. Before the fix, MX deliberately failed closed at ~5% coverage rather than activating a partial or stale model. The 18,432 unsupported entries are the model's grouped expert tensors, not arbitrary transfer failures.

Sustained across 11 consecutive refits, then re-confirmed over 12 refits, and re-confirmed again against merged ModelExpress main.

Per receiver rank per refit on the MoE arm: planned_wire_bytes 29,780,766,720 against engine_bytes 15,285,252,096, so 14,495,514,624 extra — a 1.948x amplification inherent to resharding TP2/EP4 into TP4, with 12,693,891 exact descriptors and fallback 0.

MoE Qwen3-30B-A3B: where one 10.25 s refit goes

This is a separate warm timing decomposition of the same Qwen3-30B-A3B topology, not the dense A/B above. It is the least flattering number and the most useful one for capacity planning. It is also only visible because of the stage split added in b236bc35d; before that, ModelExpress telemetry covered the install alone and reported 1.9 s of a 10.25 s refit, which looked healthy.

A warm MoE refit (Qwen3-30B-A3B) costs 10.25 s end to end:

Each phase is reported from its own median, so the rounded times and shares must not be added as if they were three spans from one trace; that is why the displayed shares total 102%. The interpretation is still stable: metadata discovery is the dominant term, publisher preparation is second, and transfer plus install is smallest.

The largest single cost in a refit is a metadata quorum check, at 2.7x the actual weight transfer. Refit is 45.6% of step time, so discover_trainers alone is roughly 22% of every training step. It re-fetches and re-parses 78,760 tensor entries across 16 trainer sources on every refit when the only thing that changed is publisher_step. The cost scales with source count rather than bytes, which is why a 30B MoE with 18,432 expert tensors exposed it and a dense 4B (~6,400 entries, ~0.4 s) did not.

This is a known, unfixed ceiling that needs an upstream ModelExpress protocol change to carry publisher_step in the list_sources record, making the check O(1) and model-size independent. Projected gain is ~10.3 s -> ~5.1 s per refit, about -22% step time.

Client-side mitigations are exhausted and measured, so please do not ask for them without new information. Fetching the per-source metadata concurrently made things 49% worse (total discover cost 5.41 s -> 8.04 s, the fetch term alone 4.02 s -> 7.56 s): the client looks idle while waiting, but the resource it waits on is the shared metadata server, so concurrency only adds contention. Skipping the shard-table rebuild did what it claimed locally, 0.79 s -> 0.20 s of parse time, but produced no end-to-end gain because the fetch dominates. eae031f60 keeps the cheap version-stamp path, and the serial fetch is pinned by a test that a concurrent implementation fails.

MoE Qwen3-30B-A3B: run-to-run variance

This is receiver-side transfer/install timing for the same MoE topology across three independently created clusters. It is not the 10.25 s framework refit total above: it isolates the receiver critical path so the source of run-to-run spread can be attributed. Fleet-critical medians were 1.845 s, 1.424 s and 1.870 s: a 1.31x across-run spread. attribute_variance.py localises it:

All of the variance is the wire; the receiver-side install is deterministic to four decimal places and has no headroom left. The straggler rank moves between runs, so this is fabric contention rather than bad placement. Practical consequence: measure any transport change with at least three cold starts, because a single run can be 1.31x off, which is larger than a plausible incremental win and can either hide a real improvement or manufacture a fake one.

What is not a performance claim

The Q=64/KV=2 TP8 -> TP8 arm described above is a correctness arm and its timings are meaningless: the model is deliberately reduced to four layers, and each TP8 group spans two 4-GPU nodes, so the transport crosses the fabric for a geometry no production run would choose. Its valid result is only: after vLLM loaded the same checkpoint, refit changed 0 of 34 parameters on all 8 receiver ranks. The same timing exclusion applies to the 12-step correctness arms, which use 512-token generation to make the gates non-vacuous rather than to benchmark latency.

Correctness gates

These are from the Qwen3-30B-A3B MoE topology, not from the dense timing arm or the reduced KV<TP model: Megatron TP2/EP4/DP2 on 16 GPUs -> four vLLM TP4 replicas on 16 GPUs. The correctness recipe runs 12 refits with 512 generated tokens and 4 prompts x 4 generations so rewards differ within a GRPO group and weights actually move. It is intentionally longer and more generation-heavy than the performance recipe.

Thresholds enforced by the functional gate, with what that 12-refit loop measured against merged ModelExpress main:

Two deliberate choices worth understanding before changing them:

• The gate asserts js_divergence_error, not an absolute gen_kl_error ceiling. Megatron and vLLM kernels have a non-zero baseline KL even when every parameter is bit-identical; measured with a refit that provably changed no parameter, that floor is 8.7e-4 to 1.3e-3 on Qwen3-30B-A3B, so a gen_kl_error < 1e-3 gate sits below its own noise floor and would fail a perfect refit. js_divergence_error is bounded and symmetric, and a single 1e-3 bound holds at every scale measured: 1.6e-4 on 0.6B, 1.3e-4 on 4B dense, 5.0e-4 on 30B MoE.
• The gate requires a non-zero gradient. Without it a recipe that produces no reward signal passes vacuously: the weights never change, so "the generator matches the trainer" is trivially true. This actually happened and invalidated three earlier readings, which is why the reward sample was also widened from 2x4 to 4x8.

Review guide

1. QKV/KV-heads-below-TP contract

Commit: 83c3d1324

Review first:

• nemo_rl/distributed/mx_megatron_helpers.py
• nemo_rl/models/policy/workers/megatron_policy_worker.py
• nemo_rl/weight_sync/mx_reshard_weight_synchronizer.py
• the matching helper, publisher, policy-worker, and synchronizer tests

Please check that layer-local config wins, local head fields are emitted only when both counts divide by TP, fused rows are validated, and only the obsolete num_query_groups % TP guard was relaxed.

2. Megatron publisher foundation

Commits: 9bcc3b8ea, c966932a0, 9446bbce4

Review parameter role classification, global expert IDs, Megatron-Bridge name order, TP/ETP ranges, and fail-closed handling of unknown or ambiguous fused layouts.

3. Weight-sync and vLLM lifecycle

Commit: 84f13c2f7, with cleanup in aa2827a29

Review:

• mx_reshard_weight_synchronizer.py
• mx_reshard_publisher.py
• mx_vllm_reshard_receiver.py
• mx_reshard_config.py
• vLLM worker/config/backend wiring

The important ordering is publish all trainer ranks, establish receiver quorum, pull/install on all receiver ranks, then commit. Any failed rank prevents version commit. Shutdown releases receivers before publishers so NIXL registrations remain valid until reads finish.

4. Reliability and observability

Commits: b236bc35d, 930f66a36, eae031f60, c473aef94

These are the commits behind the performance section. b236bc35d splits discover_s from mx_update_s and is what exposed the quorum cost; eae031f60 narrows the per-step check to the version stamp and falls back cleanly against an older ModelExpress, so the two repos need not land together; 930f66a36 makes the reporting unable to fail the refit it measures; c473aef94 adds MX_REFIT_VERIFY.

Review cleanup after exceptions, telemetry isolation, separate quorum timing, and parameter verification in mx_refit_verify.py. Telemetry must never turn a successful refit into a failed one.

5. Functional gate and portable tools

Review tests/functional/grpo_mx_reshard_refit.sh, the portable recipe configs, summarize_refit_stages.py, and attribute_variance.py.

attribute_variance.py deliberately reports stages for the rank that was fleet-critical on each step rather than the median across ranks, because those disagree: on the dense arm the all-rank median of wire_fused_s moves 1.90x between two cold starts while the fleet-critical total moves only 1.03x. Only the critical rank is on the critical path, so attributing from all-rank medians points at a stage that is not the problem.

Private namespace/image-specific cluster manifests are intentionally not included.

Evidence and scope

Validated at 32 GPUs (16 Megatron trainer GPUs, 16 vLLM receiver GPUs) on BF16:

• dense and MoE refits reach 100% coverage with zero fallback and zero unsupported;
• same-checkpoint verification finds zero changed parameters on every receiver rank;
• moving-model probability, JS-divergence, and non-zero-gradient gates pass;
• repeated refits (11, then 12) and ordered worker teardown complete cleanly;
• the performance rows above, measured with three cold starts per arm.

KV heads below TP is now proven end to end, not only unit-qualified. A reduced model with Q=64, KV=2, head_dim=128 at Megatron TP8 -> vLLM TP8 reproduced the source checkpoint bit for bit: 0 of 34 parameters changed on all 8 receiver ranks, with vLLM holding the real checkpoint before the first refit so the comparison is exact rather than statistical. Zero geometry errors, and the sparse publish pattern is visible in the per-rank tensor counts. The earlier unit and CUDA-tensor qualification (exact Q/K/V reconstruction, same-weight and changed-weight parity, projection-output equality) still holds underneath it.

This is still not full Nemotron Ultra qualification: the E2E arm is a reduced 4-layer model without EP64, so KV<TP combined with EP64 and the full Ultra architecture remain open.

Re-verified after ModelExpress #635 merged: the 12-refit MoE loop and the KV<TP exact-verify arm were both re-run against ModelExpress main at f4660d2, with every gate passing at values equal to or slightly better than the pre-merge results. That was not a formality — main also picked up four unrelated ModelExpress PRs and a refactor of the QKV alias builder itself.

FP8 installation on this path is a known gap, not a silent one: 4bbaf6a54 keeps the arm as a reproducer. The transport completes cleanly and the install then raises "Cannot copy out of meta tensor" inside ModelExpress's quantized commit path, which is a ModelExpress-side bug.

Test plan

• NeMo helper/publisher suites — 45 passed
• NeMo mx_reshard synchronizer suite — 14 passed
• policy-worker resolver wiring check
• All mx_reshard suites re-run against merged ModelExpress main — 78 passed
• ModelExpress companion reshard suites on main — 503 passed
• Real 32-GPU BF16 dense A/B and MoE performance validation, three cold starts per arm
• Real reduced Megatron TP8 -> vLLM TP8 Q=64/KV=2 E2E, bit-exact on all 8 receiver ranks
• 12-refit MoE loop re-verified against merged ModelExpress main (f4660d2)
• CI on head e45786afd
• FP8 install on this path (blocked on the ModelExpress meta-tensor bug above)
• Full Nemotron Ultra qualification, including KV<TP with EP64